import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
#1.加载模型和 Tokenizer
model_name = "deepseek-ai/deepseek-llm-7b"
tokenizer =AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
#2.加载数据集
dataset =load_dataset("json",data_files="train.json")["train"]

#3.训练参数
training_args =TrainingArguments(
    output_dir="./fine tuned model",
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    num_train_epochs=3,
    save_strategy="epoch",
    logging_dir ="./1ogs",
    report_to="none",
    fp16=True,
    optim="adamw_torch",
    evaluation_strategy="epoch",
    save_total_limit=2
)
#4.定义Trainer
trainer = Trainer(
    odel=model,
    args=training_args,
    train_datasete=dataset,
    tokenizer=tokenizer
)
#5.训练模型
trainer.train()
#6.保存训练后的模型
model.save_pretrained("fine tuned model")
tokenizer.save_pretrained("fine tuned model")